import numpy as np
import os



def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

##提取数据集的RGB像素值特征并用KNN进行分类。
from KNN import KNN
import numpy as np
import os
import time
import matplotlib.pyplot as plt

start_time = time.clock()


# 将cifar-10中的数据解压成字典类型
def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict


# 创建训练样本和测试样本
def CreatData():
    # 创建训练样本
    # 依次加载batch_data_i,并合并到x,y
    x = []
    y = []
    for i in range(1, 6):
        batch_path = 'data1\cifar-10-batches-py\data_batch_%d' % (i)
        batch_dict = unpickle(batch_path)
        train_batch = batch_dict[b'data'].astype('float')
        train_labels = np.array(batch_dict[b'labels'])
        x.append(train_batch)
        y.append(train_labels)
    # 将5个训练样本batch合并为50000x3072，标签合并为50000x1
    # np.concatenate默认axis=0，为纵向连接
    traindata = np.concatenate(x)
    trainlabels = np.concatenate(y)

    # 创建测试样本
    # 直接写cifar-10-batches-py\test_batch会报错，因此把/t当作制表符了，应用\\;
    #    test_dict=unpickle("cifar-10-batches-py\\test_batch")

    # 建议使用os.path.join()函数
    testpath = os.path.join('data1\cifar-10-batches-py', 'test_batch')
    test_dict = unpickle(testpath)
    testdata = test_dict[b'data'].astype('float')
    testlabels = np.array(test_dict[b'labels'])

    return traindata, trainlabels, testdata, testlabels


# 创建训练样本和测试样本
traindata1, trainlabels1, testdata1, testlabels1 = CreatData()
# print('traindata:',traindata.shape)
# print('trainlabels:',trainlabels.shape)
# print('testdata:',testdata.shape)
# print('testlabels:',testlabels.shape)

num_train = 10000
num_test = 5000

traindata = traindata1[:num_train]
trainlabels = trainlabels1[:num_train]
testdata = testdata1[:num_test]
testlabels = testlabels1[:num_test]

num_test = testdata.shape[0]

# K的取值
k_choice = [5, 10, 15, 20, 40,80]
# k_choice=[10]
k_accuracy = []
for k_c in k_choice:
    # 预测标签
    predictlabels = KNN(traindata, trainlabels, testdata, k=k_c)

    testlabels = np.reshape(testlabels, [num_test, 1])
    # 计算精度
    num_right = np.sum((predictlabels == testlabels).astype('float'))
    accuracy = (num_right / num_test) * 100
    k_accuracy.append(accuracy)

    print('k=%d:accuracy=%.2f%%' % (k_c, accuracy))

# 将结果可视化
plt.figure(figsize=(10, 6))
plt.plot(k_choice, k_accuracy, 'r-')
plt.plot(k_choice, k_accuracy, 'go')
plt.xlabel('k')
plt.ylabel('accuracy(%)')
plt.title('k choice')
plt.show()

# 记录最大精度以及对应K值
max_id = np.argmax(k_accuracy)
max_acc = k_accuracy[max_id]
max_acc_k = k_choice[max_id]

print('\n\nmax_acc=%.2f%%,max_acc_k=%d' % (max_acc, max_acc_k))

end_time = time.clock()

print('\n\n运行时间：%ss' % (str(end_time - start_time)))